import h5py
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.cluster import DBSCAN
import math

#FILE_NAME = 'pe_dif_sep2_98.nc' 
#FILE_NAME_CENTRAL_FORECAST = 'pe_fct_aug25_sep2.nc'
#INPUT_DATA_DIR = '/home/behollis/DATA/pierre/ocean/'
#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'

OUTPUT_DATA_DIR = '/home/brad/visweek2015-revision-code/oslines/'
INPUT_DATA_DIR = '/home/brad/DATA/ts00050/'

#realizations file 
#pe_dif_sep2_98_file = INPUT_DATA_DIR + FILE_NAME
#pe_fct_aug25_sep2_file = INPUT_DATA_DIR + FILE_NAME_CENTRAL_FORECAST 

COM =  2
LON = 53
LAT = 90
LEV = 16
DEPTH = -2.0

'''
implementation of entropy based streamline cluster from:

Chen, ChengKai, et al. An illustrative visualization framework for 
3d vector fields. Computer Graphics Forum. 
Vol. 30. No. 7. Blackwell Publishing Ltd, 2011.
'''

import numpy.linalg as la
from numpy import (array, dot, arccos, clip)
from numpy.linalg import norm

def writeElapsedTime(start, end):
    import os
    systime = os.times()
    tout = open(OUTPUT_DATA_DIR+'oceanEntropyTimings.{0}.txt'.format(systime[4]), 'w') #sys time
    tout.write('Elapsed time: {0}'.format(end - start))
    tout.close() 
 
def py_ang(v1, v2):
    ''' Returns the angle in radians between vectors 'v1' and 'v2' '''
    
    if norm(v1) == 0 or norm(v2) == 0:
        return 0.0
    
    cosang = np.dot(v1, v2) / norm(v1) / norm(v2) 
    angle = arccos(clip(cosang, -1, 1))
    
    return  angle

def calcSlineEntropies(slines):
    ''' returns avg linear and angular entropies for a seed '''
    
    aentropies = list()
    lentropies = list()
    
    for sl in slines: 
            
        total_len, sl_seg_lengths = calcSlineLength(sl) 
        
        #number of line segments
        m = len( sl_seg_lengths ) 
       
        #print 'sline lengths: ' + str(total_len)
        
        TOL_MIN = 0.000000001 
        #summation of linear entropies
        sum_lentropy = 0.0
        for s in sl_seg_lengths:
            #print s
            if s > TOL_MIN:
                sum_lentropy += ( s / total_len ) * math.log(s / total_len, 2)  
       
        lentropy = -1.0 * sum_lentropy / math.log( m + 1,2 )
       
        #find total angular variation and angle between each segment
        sl_seg_ang = list()
        total_ang = 0.0
        #total_ang_change = 0.0
        num_pts = len(sl[0])
        for j in range( 1, num_pts - 1 ):
            pt0 = ( sl[0][j-1], sl[1][j-1], 0.0 )
            pt1 = ( sl[0][j]  , sl[1][j]  , 0.0 )
            pt2 = ( sl[0][j+1], sl[1][j+1], 0.0 )
           
            v0 = np.array(pt1) - np.array(pt0)
            v1 = np.array(pt2) - np.array(pt1)
           
            ang = py_ang(v0,v1)
            #print 'ang: ' + str(ang)
           
            sl_seg_ang.append(ang)
           
            total_ang += math.fabs(ang)
            #total_ang_change += ang
           
        #summation of linear entropies
        sum_aentropy = 0.0
        for a in sl_seg_ang:
            if a != 0.0:
                sum_aentropy += ( a / total_ang ) * math.log(a / total_ang, 2) 
               
        aentropy = -1.0 * sum_aentropy / math.log( m, 2 )
       
        aentropies.append(aentropy)
        lentropies.append(lentropy)
       
    numslines = len(slines)
    sumAE = np.sum(np.asarray(aentropies))
    sumLE = np.sum(np.asarray(lentropies))
    
    return sumLE / numslines, sumAE / numslines
       
def calcSlineLength(sl):
    ''' find total length of streamline and length of each seg '''
    sl_seg_lengths = list()
    total_len = 0.0
    num_pts = len(sl[0])
    for j in range( 0, num_pts - 1 ):
        pt0 = ( sl[0][j], sl[1][j], 0.0 )
        pt1 = ( sl[0][j+1], sl[1][j+1], 0.0 )
        
        dist = np.linalg.norm( np.array(pt1) - np.array(pt0) )
        sl_seg_lengths.append( dist )
        
        total_len += dist
        
    return total_len, sl_seg_lengths
       
def calcStreamlineFeatures(slines, pts_to_sample=0.0):
    slines_features = list()
    
    for sl in slines: 
        
        if len(sl[0]) < 3 or len(sl[1]) < 3:
            continue 
        
         
        slf = list() #list of streamline feature vectors to cluster
        #find first, middle, and last points on sl
        begin_pt = [ sl[0][0], sl[1][0] ] 
        half_num_pts = len(sl[0])/2
           
        end_pt = [ sl[0][-1], sl[1][-1] ] 
        slf = begin_pt + end_pt 
           
        linear_entropy = 0.0
        angl_entropy = 0.0
           
        num_pts = len( sl[0] ) 
        num_segs = num_pts - 1
           
        total_len, sl_seg_lengths = calcSlineLength(sl)
               
        #find midpoint location of arc length
        mid_pt_idx = 0
        arc_length = 0.0
        
        for j in range( 0, num_pts - 1 ):
            pt0 = ( sl[0][j], sl[1][j], 0.0 )
            pt1 = ( sl[0][j+1], sl[1][j+1], 0.0 )
           
            dist = np.linalg.norm( np.array(pt1) - np.array(pt0) )
            arc_length += dist
          
            if arc_length >= ( total_len / 2.0 ):
                mid_pt_idx = j
                break
           
        # add mid point
        mid_pt = [ sl[ 0 ][ mid_pt_idx ], sl[ 1 ][ mid_pt_idx ] ] 
        slf += mid_pt
       
        # find additional points to be sampled
        arc_length = 0.0
        pt_count = 0.0
        if pts_to_sample > 0:
            for j in range( 0, num_pts - 1 ):
                pt0 = ( sl[0][j], sl[1][j], 0.0 )
                pt1 = ( sl[0][j+1], sl[1][j+1], 0.0 )
               
                dist = np.linalg.norm( np.array(pt1) - np.array(pt0) )
                arc_length += dist
              
                if arc_length >= pt_count * ( total_len / pts_to_sample ):
                    if pt_count >= pts_to_sample:
                        break
                   
                    pt_count += 1
                    curr_pt = [ sl[0][j], sl[1][j] ] 
                    slf += curr_pt 
                  
        slines_features.append( slf )
    
    # return feature vectors of streamlines
    return slines_features
        
def readStreamlines(x, y, f):
    slines = list()
    groups = f.keys()
    
    for mem in groups:
        dir = str(mem).zfill(4) + '/x' + str(x).zfill(3) + '/y' + str(y).zfill(3)
    
        try:
            xlst = list(f[dir][0])
            ylst = list(f[dir][1])
            zlst = 0.0 * len(list(f[dir][0]))
            
            slines.append([xlst,ylst,zlst])
        except:
            continue
    
    return slines
        
if __name__ == '__main__':
    
    import time
    start = time.time()
    
#    f = h5py.File(OUTPUT_DATA_DIR+'oceanSlinesLevel00.hdf5', 'r')
    f = h5py.File(OUTPUT_DATA_DIR+'oceanSlinesLevel00.hdf5', 'r')
    
    output = h5py.File( OUTPUT_DATA_DIR+'oceanEntropiesONLY.fixed.hdf5', 'w')
    #dslclus = output.create_dataset(name='slineclusters', shape=(X,Y), dtype='f')
    #dsampled = output.create_dataset(name='ptsSampledForClustering', shape=(X,Y), dtype='f')
    dentroL = output.create_dataset(name='avgAngularEntropy', shape=(LAT,LON), dtype='f')
    dentroA = output.create_dataset(name='avgLinearEntropy', shape=(LAT,LON), dtype='f')
    
    FEATURES = 2 * 3 + 2 * 10 #default of first, middle and end point in 2 dimensions
   
    for x in range( 0, LAT ): 
        for y in range( 0, LON ):
            print 'calculating entropy sline clusters for: {0}, {1}'.format(x, y)

            sltotal = readStreamlines(x,y,f)
            
            avg_lin_entropy = 0.0; avg_ang_entropy = 0.0
            try:
                avg_lin_entropy, avg_ang_entropy = calcSlineEntropies(sltotal)
                if math.isnan(avg_lin_entropy) :
                    avg_lin_entropy = 0.0
                    
                if math.isnan(avg_ang_entropy) :
                    avg_ang_entropy = 0.0 
            except:
                print 'error in calculating entropies...'
                
            dentroL[x,y] = avg_lin_entropy
            dentroA[x,y] = avg_ang_entropy
            print avg_ang_entropy
            print avg_lin_entropy
                  
            '''
            #calculate feature vectors for streamlines
            try:
                feat = calcStreamlineFeatures(sltotal, pts_to_sample=10)
            
                feat_total = np.array( feat )
                feat_total.reshape(len(feat_total), FEATURES)
                
                # OCEAN PARAMS
                # let's use eps to be some signficant fraction of
                # the simulation physical domain...we have 90 x 53 cells
                # five percent of the diagonal distance across the domain
                # thus, diag dist = 104.44
                # 0.05 * 104.44 = 5.2
                # NOTE: include this metric and value in the paper
                # min samples is 0.1 of total members, or 3 
                
                # STIRRING PARAMS
                #dcov[x,y] = np.cov(np.array(tpts).T)#pca.get_covariance()
                # let's use eps to be some signficant fraction of
                # the simulation physical domain...we have 152 x 152 cells
                # five percent of the diagonal distance across the domain
                # thus, diag dist = 214.96
                # 0.05 * 214.96 = 10.7
                # NOTE: include this metric and value in the paper
                # min samples is 0.1 of total members, or 2 
                
                db = DBSCAN(eps=10.7, min_samples = 2).fit( np.asarray(feat_total) )
                
                core_samples_mask = np.zeros_like(db.labels_, dtype=bool)
                core_samples_mask[ db.core_sample_indices_] = True
                labels = db.labels_
        
                # Number of clusters in labels, ignoring noise if present.
                n_clusters_ = len(set(labels)) - (1 if -1 in labels else 0)
                dslclus[x,y] = n_clusters_
                print n_clusters_
            except:
                dslclus[x,y] = 0
                
            dsampled[x,y] = FEATURES
            '''
            #print 'clusters: ' + str(dslclus[x,y])
            
    end = time.time()
    writeElapsedTime(start, end)

    f.close()
    output.close()
    
    print 'finished!'